'''
Author: caishuyang
Date: 2023-03-13 17:00:32
LastEditors: caishuyang
LastEditTime: 2023-03-13 21:16:58
Description: 测试环节
'''
import cv2
import torch
import os
import numpy as np
from PIL import Image
from _03Deeplabv3plus import DeepLabV3
from torchvision.transforms import transforms

modelpath=os.path.join("Model","model_01.pth")
videopath=os.path.join("Video","test.mp4")

deeplab=DeepLabV3()
deeplab.load_state_dict(torch.load(modelpath))
deeplab.eval()
trans=transforms.ToTensor()
video=cv2.VideoCapture(videopath)
flag=True

while flag:
    flag,frame=video.read()
    if frame is None:
        break
    img=cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
    img=Image.fromarray(img)
    img=img.resize((160, 120),Image.ANTIALIAS)
    img=trans(img)
    img=np.array(img)
    img=np.array([img])
    road=deeplab(torch.tensor(img))
    roadimg=road[0,0,:,:].detach().numpy()
    back=road[0,1,:,:].detach().numpy()
    roadimg=(roadimg-back)+1
    roadimg=cv2.resize(roadimg,(640,480))
    roadimg=np.array(70*roadimg,dtype="uint8")
    frame=0.5*frame
    frame[:,:,1]+=0.5*roadimg
    cv2.imshow("show",frame)
    cv2.waitKey(1)

